package cn.iocoder.yudao.module.liveprice.websocket;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

/**
 * WebSocket连接管理器
 */
@Slf4j
@Component
public class PriceWebSocketManager {

    /**
     * 按主题分类的会话集合
     */
    private final Map<String, CopyOnWriteArraySet<WebSocketSession>> topicSessionsMap = new ConcurrentHashMap<>();

    /**
     * 添加会话到特定主题
     *
     * @param topic   主题名称
     * @param session WebSocket会话
     */
    public void addSession(String topic, WebSocketSession session) {
        topicSessionsMap.computeIfAbsent(topic, k -> new CopyOnWriteArraySet<>()).add(session);
        log.info("添加会话到主题: {}, 会话ID: {}", topic, session.getId());
    }

    /**
     * 从特定主题移除会话
     *
     * @param topic   主题名称
     * @param session WebSocket会话
     */
    public void removeSession(String topic, WebSocketSession session) {
        CopyOnWriteArraySet<WebSocketSession> sessions = topicSessionsMap.get(topic);
        if (sessions != null) {
            sessions.remove(session);
            log.info("从主题移除会话: {}, 会话ID: {}", topic, session.getId());
        }
    }

    /**
     * 向特定主题的所有会话发送消息
     *
     * @param topic   主题名称
     * @param message 消息内容
     */
    public void sendMessageToTopic(String topic, String message) {
        CopyOnWriteArraySet<WebSocketSession> sessions = topicSessionsMap.get(topic);
        if (sessions != null && !sessions.isEmpty()) {
            TextMessage textMessage = new TextMessage(message);
            for (WebSocketSession session : sessions) {
                try {
                    if (session.isOpen()) {
                        session.sendMessage(textMessage);
                    }
                } catch (IOException e) {
                    log.error("发送消息到会话失败: {}", session.getId(), e);
                }
            }
            log.info("向主题[{}]的{}个会话发送消息", topic, sessions.size());
        }
    }
} 